package com.epam.training.collection;

import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	private static final boolean RED = false;
	private static final boolean BLACK = true;

	private SimpleEntry<K, V> root = null;
	private Comparator<? super K> comparator;

	private int size = 0;

	MyTreeMap() {
		comparator = null;
	}

	MyTreeMap(Comparator<? super K> comparator) {
		this.comparator = comparator;
	}

	@Override
	public void clear() {
		size = 0;
		root = null;
	}

	@Override
	public boolean containsKey(K key) {
		return (getEntry(key) != null) ? true : false;
	}

	private SimpleEntry<K, V> getFirstEntry() {
		SimpleEntry<K, V> p = root;
		if (p != null) {
			while (p.left != null) {
				p = p.left;
			}
		}
		return p;
	}

	@Override
	public boolean containValue(V value) {
		for (SimpleEntry<K, V> e = getFirstEntry(); e != null; e = successor(e)) {
			if (value.equals(e.value)) {
				return true;
			}
		}

		return false;
	}

	@Override
	public V get(K key) {
		SimpleEntry<K, V> p = getEntry(key);
		return (p == null ? null : p.value);
	}

	@Override
	public boolean isEmpty() {
		return size == 0;
	}

	@SuppressWarnings("unchecked")
	private SimpleEntry<K, V> getEntryUsingComparator(Object key) {
		K k = (K) key;
		Comparator<? super K> cpr = comparator;
		if (cpr != null) {
			SimpleEntry<K, V> p = root;
			while (p != null) {
				int cmp = cpr.compare(k, p.key);
				if (cmp < 0) {
					p = p.left;
				} else if (cmp > 0) {
					p = p.right;
				} else {
					return p;
				}
			}
		}
		return null;
	}

	@SuppressWarnings("unchecked")
	SimpleEntry<K, V> getEntry(Object key) {
		if (comparator != null) {
			return getEntryUsingComparator(key);
		}
		if (key == null) {
			throw new NullPointerException();
		}

		Comparable<? super K> k = (Comparable<? super K>) key;
		SimpleEntry<K, V> p = root;
		while (p != null) {
			int cmp = k.compareTo(p.key);
			if (cmp < 0) {
				p = p.left;
			} else if (cmp > 0) {
				p = p.right;
			} else {
				return p;
			}
		}
		return null;
	}

	@SuppressWarnings("unchecked")
	@Override
	public V put(K key, V value) {
		SimpleEntry<K, V> t = root;

		if (t == null) {
			root = new SimpleEntry<K, V>(key, value, null);

			size = 1;
			return null;
		}

		int cmp;
		SimpleEntry<K, V> parent;

		Comparator<? super K> cpr = comparator;
		if (cpr != null) {
			do {
				parent = t;
				cmp = cpr.compare(key, t.key);
				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					return t.setValue(value);
				}
			} while (t != null);
		} else {
			if (key == null) {
				throw new NullPointerException();
			}

			Comparable<? super K> k = (Comparable<? super K>) key;
			do {
				parent = t;
				cmp = k.compareTo(t.key);
				if (cmp < 0) {
					t = t.left;
				} else if (cmp > 0) {
					t = t.right;
				} else {
					return t.setValue(value);
				}
			} while (t != null);
		}
		SimpleEntry<K, V> e = new SimpleEntry<>(key, value, parent);
		if (cmp < 0) {
			parent.left = e;
		} else {
			parent.right = e;
		}
		insertFixup(e);
		size++;
		return null;
	}

	private void insertFixup(SimpleEntry<K, V> node) {

		SimpleEntry<K, V> y;

		while (node.parent.color == RED) {
			if (node.parent == node.parent.parent.left) {
				y = node.parent.parent.right;

				if (y.color == RED) {
					node.parent.color = BLACK;
					y.color = BLACK;
					node.parent.parent.color = RED;
					node = node.parent.parent;
				} else if (node.equals(node.parent.right)) {
					node = node.parent;
					rotateLeft(node);
				}
				node.parent.color = BLACK;
				node.parent.parent.color = RED;
				rotateRight(node);
			} else if (node.parent == node.parent.parent.right) {
				y = node.parent.parent.left;

				if (y.color == RED) {
					node.parent.color = BLACK;
					y.color = BLACK;
					node.parent.parent.color = RED;
					node = node.parent.parent;
				} else if (node.equals(node.parent.left)) {
					node = node.parent;
					rotateRight(node);
				}
				node.parent.color = BLACK;
				node.parent.parent.color = RED;
				rotateLeft(node);
			}
		}

		root.color = BLACK;
	}

	private void rotateLeft(SimpleEntry<K, V> p) {
		if (p.right == null) {
			return;
		}

		SimpleEntry<K, V> right = p.right;
		p.right = right.left;

		if (right.left != null) {
			right.left.parent = p;
		}

		right.parent = p.parent;

		if (p.parent == null) {
			root = right;
		} else if (p == p.parent.left) {
			p.parent.left = right;
		} else {
			p.parent.right = right;
		}

		right.left = p;
		p.parent = right;
	}

	private void rotateRight(SimpleEntry<K, V> p) {
		if (p.left == null) {
			return;
		}

		SimpleEntry<K, V> left = p.left;
		p.left = left.right;

		if (left.right != null) {
			left.right.parent = p;
		}

		left.parent = p.parent;

		if (p.parent == null) {
			root = left;
		} else if (p == p.parent.right) {
			p.parent.right = left;
		} else {
			p.parent.left = left;
		}

		left.right = p;
		p.parent = left;
	}

	SimpleEntry<K, V> successor(SimpleEntry<K, V> t) {
		if (t == null) {
			return null;
		} else if (t.right != null) {
			SimpleEntry<K, V> p = t.right;
			while (p.left != null) {
				p = p.left;
			}
			return p;
		} else {
			SimpleEntry<K, V> p = t.parent;
			SimpleEntry<K, V> ch = t;
			while (p != null && ch == p.right) {
				ch = p;
				p = p.parent;
			}
			return p;
		}
	}

	@Override
	public V remove(K key) {
		SimpleEntry<K, V> p = getEntry(key);

		if (p == null) {
			return null;
		}

		size--;
		V oldValue = p.value;

		if (p.left != null && p.right != null) {
			SimpleEntry<K, V> s = successor(p);
			p.key = s.key;
			p.value = s.value;
			p = s;
		}

		SimpleEntry<K, V> replacement = (p.left != null ? p.left : p.right);

		if (replacement != null) {
			replacement.parent = p.parent;
			if (p.parent == null) {
				root = replacement;
			} else if (p == p.parent.left) {
				p.parent.left = replacement;
			} else {
				p.parent.right = replacement;
			}

			p.left = p.right = p.parent = null;

			if (p.color == BLACK)
				deleteFixup(replacement);
		} else if (p.parent == null) {
			root = null;
		} else {
			if (p.color == BLACK) {
				deleteFixup(p);
			}

			if (p.parent != null) {
				if (p == p.parent.left) {
					p.parent.left = null;
				} else if (p == p.parent.right) {
					p.parent.right = null;
				}

				p.parent = null;
			}
		}

		return oldValue;
	}

	private void deleteFixup(SimpleEntry<K, V> node) {
		while (!node.equals(root) && node.color == BLACK) {
			if (node.equals(node.parent.left)) {
				SimpleEntry<K, V> w = node.parent.right;
				if (w.color == RED) {
					w.color = BLACK;
					node.parent.color = RED;
					rotateLeft(node.parent);
					w = node.parent.right;
				}

				if (w.left.color == BLACK && w.right.color == BLACK) {
					w.color = RED;
					node = node.parent;
				}

				else if (w.right.color == BLACK) {
					w.left.color = BLACK;
					w.color = RED;
					rotateRight(w);
					w = node.parent.right;
				}

				w.color = node.parent.color;
				node.parent.color = BLACK;
				w.right.color = BLACK;
				rotateLeft(node.parent);
				node = root;

			} else {
				if (node.equals(node.parent.right)) {
					SimpleEntry<K, V> w = node.parent.left;
					if (w.color == RED) {
						w.color = BLACK;
						node.parent.color = RED;
						rotateRight(node.parent);
						w = node.parent.left;
					}

					if (w.right.color == BLACK && w.left.color == BLACK) {
						w.color = RED;
						node = node.parent;
					}

					else if (w.left.color == BLACK) {
						w.right.color = BLACK;
						w.color = RED;
						rotateLeft(w);
						w = node.parent.left;
					}

					w.color = node.parent.color;
					node.parent.color = BLACK;
					w.left.color = BLACK;
					rotateRight(node.parent);
					node = root;
				}
			}
		}
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public Iterator<MyMap.Entry<K, V>> entryIterator() {
		return new EntryIterator(getFirstEntry());
	}

	class EntryIterator implements Iterator<MyMap.Entry<K, V>> {
		SimpleEntry<K, V> next;
		SimpleEntry<K, V> lastReturned;

		EntryIterator(SimpleEntry<K, V> first) {
			lastReturned = null;
			next = first;
		}

		@Override
		public boolean hasNext() {
			return next != null;
		}

		@Override
		public MyMap.Entry<K, V> next() {
			SimpleEntry<K, V> e = next;

			if (e == null) {
				throw new NoSuchElementException();
			}

			next = successor(e);
			lastReturned = e;

			return e;
		}

		@Override
		public void remove() {
			if (lastReturned == null)
				throw new IllegalStateException();

			if (lastReturned.left != null && lastReturned.right != null) {
				next = lastReturned;
			}

			MyTreeMap.this.remove(lastReturned.key);
			lastReturned = null;
		}
	}

	@SuppressWarnings("hiding")
	class SimpleEntry<K, V> implements Entry<K, V> {
		V value;
		K key;

		SimpleEntry<K, V> left = null;
		SimpleEntry<K, V> right = null;

		SimpleEntry<K, V> parent = null;

		boolean color = BLACK;

		SimpleEntry(K key, V value, SimpleEntry<K, V> parent) {
			this.key = key;
			this.value = value;
			this.parent = parent;
		}

		@Override
		public K getKey() {
			return key;
		}

		@Override
		public V getValue() {
			return value;
		}

		@Override
		public V setValue(V value) {
			final V old = this.value;
			this.value = value;

			return old;
		}

		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + getOuterType().hashCode();
			result = prime * result + (color ? 1231 : 1237);
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj) {
				return true;
			}
			if (obj == null) {
				return false;
			}
			if (getClass() != obj.getClass()) {
				return false;
			}
			@SuppressWarnings("unchecked")
			SimpleEntry<K, V> other = (SimpleEntry<K, V>) obj;
			if (!getOuterType().equals(other.getOuterType())) {
				return false;
			}
			if (color != other.color) {
				return false;
			}
			return true;
		}

		@SuppressWarnings("unchecked")
		private MyTreeMap<K, V> getOuterType() {
			return (MyTreeMap<K, V>) MyTreeMap.this;
		}

	}

}
